import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import LineInspector, PanTool, ZoomTool
from enable.api import ComponentEditor

class Form(HasTraits):
    alpha = Range(np.pi, 2*np.pi, 4.02595098557532)
    plot = Instance(Plot)

    factor = Enum('cot1,cot2', 'X,Y', 'f,g')
    diffraction = Enum('GTD', 'PTD')

    def traits_init(self):
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            phi=[],
            y1=[],
            y2=[],
        )
        self.plot.plot(("phi", "y1"), color='auto')
        self.plot.plot(("phi", "y2"), color='auto')
        self.update_data()

    @observe('alpha,factor,diffraction')
    def update_data(self, event=None):
        phi = np.linspace(0, self.alpha, 201)
        n = self.alpha / np.pi
        match self.diffraction:
            case 'GTD':
                cot1 = -1 / 2 / n / np.tan((np.pi - 2 * phi) / 2 / n)
                cot2 = -1 / 2 / n / np.tan((np.pi + 2 * phi) / 2 / n)
                cot1[np.isinf(cot1)]=np.nan
                cot2[np.isinf(cot2)]=np.nan
                X = np.full_like(phi,-1 / n / np.tan(np.pi / 2 / n))
                Y = cot1+cot2
                f = X - Y
                g = X + Y
                match self.factor:
                    case 'cot1,cot2':
                        y1=cot1
                        y2=cot2
                    case 'X,Y':
                        y1 = X
                        y2 = Y
                    case 'f,g':
                        y1 = f
                        y2 = g
            case 'PTD':
                cot1 = -1 / 2 / n / np.tan((np.pi - 2 * phi) / 2 / n)+1/2/np.tan((np.pi - 2 * phi) / 2)*(phi<np.pi)
                cot2 = -1 / 2 / n / np.tan((np.pi + 2 * phi) / 2 / n)+1/2/np.tan((np.pi - 2 * (self.alpha-phi)) / 2)*(self.alpha-phi<np.pi)
                cot1[np.abs(phi-np.pi/2)<.001]=0
                cot2[np.abs(self.alpha-phi-np.pi/2)<.001]=0
                X = np.full_like(phi,-1 / n / np.tan(np.pi / 2 / n))
                Y = cot1+cot2
                f = X - Y
                g = X + Y
                match self.factor:
                    case 'cot1,cot2':
                        y1=cot1
                        y2=cot2
                    case 'X,Y':
                        y1 = X
                        y2 = Y
                    case 'f,g':
                        y1 = f
                        y2 = g
        self.plot.data.update_data(
            phi=phi,
            y1=y1,
            y2=y2,
        )

    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        Item('alpha', label='尖劈角度', editor_args={'mode': 'slider'}),
        HGroup(
            Item('factor', label='函数类型', style='custom', editor_args={'cols': 8}),
            spring,
        ),
        HGroup(
            Item('diffraction', label='绕射系数', style='custom', editor_args={'cols': 8}),
            spring,
        ),
        width=1400,
        height=1000,
        resizable=True,
    )


if __name__ == '__main__':
    form = Form()
    form.configure_traits()
